{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-13T03:13:27.183972Z",
     "start_time": "2024-12-13T03:13:23.017327Z"
    }
   },
   "source": [
    "import os\n",
    "import json\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "from ResNet import resnet50\n",
    "import matplotlib.pyplot as plt"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "classarr = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'Random', 'Scratch']",
   "id": "21098ea738bce0b1"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "def plot_confusion_matrix(cm, savename, title='Confusion Matrix'):\n",
    "    plt.figure(figsize=(12, 8), dpi=100)\n",
    "    np.set_printoptions(precision=2)\n",
    "\n",
    "    # 在混淆矩阵中每格的概率值\n",
    "    ind_array = np.arange(len(classarr))\n",
    "    x, y = np.meshgrid(ind_array, ind_array)\n",
    "    for x_val, y_val in zip(x.flatten(), y.flatten()):\n",
    "        c = cm[y_val][x_val]\n",
    "        if c > 0.001:\n",
    "            plt.text(x_val, y_val, \"%d\" % (c,), color='red', fontsize=15, va='center', ha='center')\n",
    "\n",
    "    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)\n",
    "    plt.title(title)\n",
    "    plt.colorbar()\n",
    "    xlocations = np.array(range(len(classarr)))\n",
    "    plt.xticks(xlocations, classarr, rotation=90)\n",
    "    plt.yticks(xlocations, classarr)\n",
    "    plt.ylabel('Actual label')\n",
    "    plt.xlabel('Predict label')\n",
    "\n",
    "    # offset the tick\n",
    "    tick_marks = np.array(range(len(classarr))) + 0.5\n",
    "    plt.gca().set_xticks(tick_marks, minor=True)\n",
    "    plt.gca().set_yticks(tick_marks, minor=True)\n",
    "    plt.gca().xaxis.set_ticks_position('none')\n",
    "    plt.gca().yaxis.set_ticks_position('none')\n",
    "    plt.grid(True, which='minor', linestyle='-')\n",
    "    plt.gcf().subplots_adjust(bottom=0.15)\n",
    "\n",
    "    # show confusion matrix\n",
    "    plt.savefig(savename, format='png')\n",
    "    plt.show()\n"
   ],
   "id": "16c3e4ac0c00f543"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.Resize((26, 26)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])\n",
    "\n",
    "# load image\n",
    "# 指向需要遍历预测的图像文件夹\n",
    "class_names = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'Random', 'Scratch']\n",
    "img_path_list = []\n",
    "for i in range(len(class_names)):\n",
    "    imgs_root = \"C://Users//w8887757//PycharmProjects//wm811k//data//test//\" + class_names[i]\n",
    "    img_path_list += [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(\".jpg\")]\n",
    "\n",
    "# imgs_root = \"./input/val\"\n",
    "# assert os.path.exists(imgs_root), f\"file: '{imgs_root}' dose not exist.\"\n",
    "# # 读取指定文件夹下所有jpg图像路径\n",
    "# img_path_list = [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(\".jpg\")]\n",
    "\n",
    "# read class_indict\n",
    "class_indict = {'0':'Center','1':'Donut','2':'Edge-Loc','3':'Edge-Ring','4':'Loc','5':'Near-full','6':'Random','7':'Scratch'}\n",
    "\n",
    "# create model\n",
    "model = resnet50(num_classes=8).to(device)\n",
    "\n",
    "# load model weights\n",
    "weights_path = \"./input/save_weights/resnet50_dataAug_epo150_class8.pth\"\n",
    "assert os.path.exists(weights_path), f\"file: '{weights_path}' dose not exist.\"\n",
    "model.load_state_dict(torch.load(weights_path, map_location=device))\n",
    "\n",
    "# prediction\n",
    "model.eval()\n",
    "batch_size = 8  # 每次预测时将多少张图片打包成一个batch\n",
    "y_true = []\n",
    "y_pred = []\n",
    "truecnt = 0;\n",
    "with torch.no_grad():\n",
    "    for ids in range(0, len(img_path_list) // batch_size):\n",
    "        img_list = []\n",
    "        for img_path in img_path_list[ids * batch_size: (ids + 1) * batch_size]:\n",
    "            assert os.path.exists(img_path), f\"file: '{img_path}' dose not exist.\"\n",
    "            img = Image.open(img_path)\n",
    "            img = data_transform(img)\n",
    "            img_list.append(img)\n",
    "\n",
    "        # batch img\n",
    "        # 将img_list列表中的所有图像打包成一个batch\n",
    "        batch_img = torch.stack(img_list, dim=0)\n",
    "        # predict class\n",
    "        output = model(batch_img.to(device)).cpu()\n",
    "        predict = torch.softmax(output, dim=1)\n",
    "        probs, classes = torch.max(predict, dim=1)\n",
    "\n",
    "        for idx, (pro, cla) in enumerate(zip(probs, classes)):\n",
    "            print(str(cla.numpy()))\n",
    "            y_pred.append(class_indict[str(cla.numpy())])\n",
    "            y_true.append(str(img_path_list[ids * batch_size + idx])[len(\"C://Users//w8887757//PycharmProjects//wm811k//data//test//\"):].split(\"\\\\\")[0])\n",
    "            if (class_indict[str(cla.numpy())] == str(img_path_list[ids * batch_size + idx])[len(\"C://Users//w8887757//PycharmProjects//wm811k//data//test//\"):].split(\"\\\\\")[0]):\n",
    "                truecnt = truecnt + 1\n",
    "            print(\"image: {}  class: {}  prob: {:.3}\".format(img_path_list[ids * batch_size + idx],\n",
    "                                                             class_indict[str(cla.numpy())],\n",
    "                                                             pro.numpy()))\n",
    "predict_acc = truecnt /len(y_true)\n",
    "print(\"准确率：\" + str(predict_acc))\n",
    "cm = confusion_matrix(y_true, y_pred)\n",
    "plot_confusion_matrix(cm, 'confusion_matrix_data8_epo100.png', title='confusion matrix')"
   ],
   "id": "b7cd543490746c62"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
